"""
=============================================================
        LangChain 异步流式输出 + JSON解析示例
=============================================================

功能：
    - 使用 ChatOpenAI 调用大模型生成内容。
    - 构建 ChatPromptTemplate 提示词。
    - 使用 JsonOutputParser 将模型输出解析为 JSON。
    - 异步流式输出大模型生成结果，可边生成边处理。

用途：
    - 大模型生成结构化 JSON 数据。
    - 异步流式输出，适合实时显示模型生成内容。
"""

# -------------------------------
# 导入依赖
# -------------------------------
import asyncio
import os

from dotenv import load_dotenv
from langchain_core.output_parsers import JsonOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI

# -------------------------------
# 1️⃣ 自动加载 .env 文件中的环境变量
# -------------------------------
load_dotenv()

# -------------------------------
# 2️⃣ 初始化大模型
# -------------------------------
llm = ChatOpenAI(api_key=os.getenv("XUNFEI_API_KEY"), openai_api_base=os.getenv("XUNFEI_API_BASE"), model="xdeepseekv32exp", streaming=True, temperature=0.7)

# -------------------------------
# 3️⃣ 构建提示词
# -------------------------------
prompt = ChatPromptTemplate.from_template(
    "请以JSON格式输出中国每个省份的人口列表，"
    "要求使用一个外层键为'sheng'的字典，里面是所有省份的数据列表，"
    "每个省份的对象必须包含键'name'（省份名称）和'value'（人口数）。"
)

# -------------------------------
# 4️⃣ 输出解析器：将模型输出解析为 JSON
# -------------------------------
parser = JsonOutputParser()

# -------------------------------
# 5️⃣ 组合 chain：提示词 | 模型 | 输出解析器
# -------------------------------
chain = prompt | llm | parser


# -------------------------------
# 6️⃣ 异步流式输出函数
# -------------------------------
async def async_stream():
    """
    通过 chain.astream 异步获取生成内容，每次得到模型生成的一小块结果。
    可边生成边处理或展示。
    """
    async for chunk in chain.astream({}):
        print("\n")
        print(chunk, end="", flush=True)


# -------------------------------
# 7️⃣ 运行异步函数
# -------------------------------
asyncio.run(async_stream())
